from khronos import *
from khronos.optimization import *
from khronos.extras.components import *
from random import Random

# ------------------------------------------------------------------------------
# ENTITY DESCRIPTION
# ------------------------------------------------------------------------------
class Customer(Atomic):
    wait = 0 * SECOND
    state = "<uninitialized>"
    def status(self):
        return self.state
        
    @chain
    def go(self, gas_station, fill_time, pay_time, max_wait):
        pumps    = gas_station.find("pumps")
        cashiers = gas_station.find("cashiers")
        
        fill = yield self.fill(pumps, fill_time, max_wait)
        if fill.activated():
            yield self.pay(cashiers, pay_time)
            
        self.parent.observe(self.wait)
        self.terminate()
        
    @chain
    def fill(self, pumps, fill_time, max_wait):
        pump = pumps.get_free()
        request = pump.acquire()
        self.state = "waiting for pump %s" % (pump.name,)
        
        # Request a free pump (or the one with the smallest queue).
        start = self.get_timedelta()
        yield request % max_wait
        self.wait = self.get_timedelta() - start
        
        # If a pump wasn't acquired within the max_wait period, cancel this chain.
        if request.canceled():
            raise ChainCancel()
            
        # Pump acquired, time to fill up.
        self.state = "filling at pump %s" % (pump.name,)
        yield fill_time
        # Do not forget to release the pump for the other customers!
        pump.release()
        
    @chain
    def pay(self, cashiers, pay_time):
        cashier = cashiers.get_free()
        self.state = "waiting for cashier %s" % (cashier.name,)
        
        start = self.get_timedelta()
        yield cashier.acquire()
        self.wait += self.get_timedelta() - start
        
        self.state = "paying at cashier %s" % (cashier.name,)
        yield pay_time
        cashier.release()
        
class Customers(Composite):
    mandatory_config = ["gas_station", "arrival_rate", "fill_time", "pay_time", "max_wait"]
    @chain
    def initializer(self):
        self.clear()
        self.validate_config()
        self.config.arrival_rate.set_clock(self.get_clock())
        gas_station = self.find(self.config.gas_station)
        # Initialize statistics.
        self.wait_times = TSeries(storing=True, timescale=HOUR)
        n = 0
        while True:
            yield self.config.arrival_rate()
            customer = Customer(name=str(n), parent=self)
            self.launch(customer, initializer=Call(customer.go, gas_station, 
                                                   self.config.fill_time(), 
                                                   self.config.pay_time(), 
                                                   self.config.max_wait()))
            n += 1
            
    def observe(self, wait):
        # Register the wait in minutes.
        self.wait_times.register(wait / MINUTE, self.get_timedelta())
        
#    def finalizer(self):
#        gas_station = self.find(self.config.gas_station)
#        npumps = gas_station.config.pumps
#        ncashiers = gas_station.config.cashiers
#        
#        p = Plotter("Gas station example", rows=2, cols=1)
#        avg_wait = self.wait_times.mean_series()
#        avg_wait.run_chart(title="Pumps=%d, Cashiers=%d" % (npumps, ncashiers),
#                           xlabel="Time (h)", ylabel="Average wait (m)", 
#                           show_mean=False, plotter=p)
#        self.wait_times.run_chart(title="Wait times", show_mean=False, plotter=p)
        
class GasStation(Composite):
    mandatory_config = ["pumps", "cashiers"]
    def initializer(self):
        self.validate_config()
        self.clear()
        self.add(ResourcePool("pumps",    config=self.config.pumps), 
                 ResourcePool("cashiers", config=self.config.cashiers))
                 
class ResourcePool(Composite):
    def initializer(self):
        self.clear()
        for x in xrange(self.config):
            self.add(Resource(name=str(x)))
            
    def get_free(self):
        min_queue = INF
        max_avail = 0
        best_pumps = None
        for rsc in self:
            rsc_queue = rsc.queue_size()
            rsc_avail = rsc.get_available()
            if rsc_queue < min_queue or \
            (rsc_queue == min_queue and rsc_avail > max_avail):
                min_queue = rsc_queue
                max_avail = rsc_avail
                best_pumps = [rsc]
            elif rsc_queue == min_queue and rsc_avail == max_avail:
                best_pumps.append(rsc)
        return self.get_rng().choice(best_pumps)
        
# ------------------------------------------------------------------------------
# MODEL BUILDING
# ------------------------------------------------------------------------------
def model():
    rng = Random()
    
    def h_rate(): return max(0.0, rng.expovariate(60.0)) * HOUR
    def m_rate(): return max(0.0, rng.expovariate(30.0)) * HOUR
    def l_rate(): return max(0.0, rng.expovariate(10.0)) * HOUR
    arrival_rate = Distribution(default=l_rate, loop=DAY, name="arrival_rate")
    arrival_rate.add(h_rate, 
                     HOUR.multiply(8.0, 8.5))
    arrival_rate.add(m_rate, 
                     HOUR.multiply(7.5, 8.0), 
                     HOUR.multiply(8.5, 9.0), 
                     HOUR.multiply(18.0, 20.0))
    
    def fill_time(): return max(0.0, rng.gauss(4.0, 1.0)) * MINUTE
    def pay_time():  return max(0.0, rng.gauss(1.0, 0.2)) * MINUTE
    def max_wait():  return max(0.0, rng.gauss(10.0, 2.0)) * MINUTE
    
    m = Composite("root", config=rng)
    GasStation("gas_station", parent=m, config=dict(pumps=4, cashiers=2))
    Customers("customers",    parent=m, config=dict(gas_station=".gas_station", 
                                                    arrival_rate=arrival_rate, 
                                                    fill_time=fill_time, 
                                                    pay_time=pay_time, 
                                                    max_wait=max_wait))
    return m
    
def main():
    m = model()
    k = Khronos(m, rng=m.config)
    k.single_run(WEEK, seed=1, trace=False)
    
# ------------------------------------------------------------------------------
# EXPERIMENT DESCRIPTION
# ------------------------------------------------------------------------------
m = model()
k = Khronos(m, rng=m.config)

class GasExperiment(Experiment):
    runs = 10
    duration = WEEK
    indicators = Namespace(meanwait=StatIndicator.subclass("MeanWait"))
    
    def at_end(self):
        cfg = self.model.find("gas_station").config
        p = cfg.pumps
        c = cfg.cashiers
        meanwait = self.indicators.meanwait.stat
        meanwait.run_chart(title="Pumps=%d, Cashiers=%d" % (p, c), show_stddev=False)
        
@GasExperiment.indicators.meanwait.getter
def run_meanwait(indicator):
    model = indicator.experiment.model
    return model.find("customers").wait_times.mean()
    
# ------------------------------------------------------------------------------
# OPTIMIZATION PROBLEM DESCRIPTION
# ------------------------------------------------------------------------------
pumps    = Variable("pumps",    Domain([1, 2, 3, 4, 5]))
cashiers = Variable("cashiers", Domain([1, 2, 3]))
prob     = Problem()
prob.add_variables([pumps, cashiers])

@pumps.setter
def set_pumps(model, n):
    model.find("gas_station").config["pumps"] = n
    
@cashiers.setter
def set_cashiers(model, n):
    model.find("gas_station").config["cashiers"] = n
    
def exp_meanwait(experiment):
    return experiment.indicators.meanwait.stat.wmean()
    
opt = optimizer.bruteforce()
#opt.solve(prob, objective=exp_meanwait, model=m, experiment=GasExperiment, simulator=k)